/************************************************************************
Copyright 2003, 2004, 2005, 2006, 2007, 2008 Grid Technology Research Center,
National Institute of Advanced Industrial Science and Technology.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
************************************************************************/

package jp.aist.gtrc.plus.scheduler.specific.sge.sched;

import java.io.IOException;
import java.io.InputStream;
import java.net.Socket;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.Locale;

import jp.aist.gtrc.plus.scheduler.node.NodeInfo;
import jp.aist.gtrc.plus.scheduler.specific.sge.jaxb.common.*;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEHostGroup;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEJobID;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEJobStatus;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGENodeStatus;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEParallelEnv;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEProject;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEQInstance;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEQueueStatus;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEServerStatus;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGETaskStatus;
import jp.aist.gtrc.plus.scheduler.specific.sge.status.SGEUserset;
import jp.aist.gtrc.plus.scheduler.status.JobID;
import jp.aist.gtrc.plus.scheduler.status.JobStatus;
import jp.aist.gtrc.plus.scheduler.status.MainServer;
import jp.aist.gtrc.plus.scheduler.status.NodeStatus;
import jp.aist.gtrc.plus.scheduler.status.QueueStatus;
import jp.aist.gtrc.plus.scheduler.status.ServerStatus;
import jp.aist.gtrc.plus.scheduler.sys.SchedulerException;
import jp.aist.gtrc.plus.scheduler.sys.SignalType;
import jp.aist.gtrc.plus.scheduler.util.TextUtil;

import org.apache.log4j.Logger;

public final class SGESelfMainServer implements MainServer {

	private Socket socket;
	private static OpedXMLNormalizer xmlNormalizer;
	private static SGEOperatord operatord;
	private static SGESelfMainServer srvInstance;

	protected static final  Logger logger = Logger.getLogger(SGESelfMainServer.class);

	private SGESelfMainServer() throws SchedulerException {

		if (operatord == null) {
			operatord = new SGEOperatord();
		}
		if (xmlNormalizer == null) {
			xmlNormalizer = new OpedXMLNormalizer();
		}

		Runtime.getRuntime().addShutdownHook(new Thread() {
		    public void run() {
		    	if (srvInstance != null) {
		    		srvInstance.shutdown();
		    		srvInstance = null;
		    	}
		    }
		});
	}

	public static synchronized SGESelfMainServer getInstance() throws SchedulerException {

		if (srvInstance == null) {
			srvInstance = new SGESelfMainServer();
		}
		srvInstance.connect();
		return srvInstance;

	}

	private void connect() throws SchedulerException {

		socket = operatord.connect();

	}

	public boolean isConnected() {

		return operatord.isConnected();

	}

	public void disconnect() {

		operatord.disconnect();

	}

	public void shutdown() {

		operatord.shutdown();
	}

	private Object getStatus(String type, Class<?> desiredClass) throws SchedulerException {

		operatord.sendRequest("GetStatus " + type);

		InputStream stream;
		try {
			stream = socket.getInputStream();
		} catch (IOException e) {
			throw new SchedulerException(e.getMessage());
		}

		Object o =  xmlNormalizer.makeJAXBObject(stream, type);
		if ((o != null) && o.getClass().equals(desiredClass)) {
			return o;
		} else {
			throw new SchedulerException("Invalid JAXB class object was made");
		}

	}

	public ServerStatus getServerStatus() throws SchedulerException {

		JAXBSGEServerList jaxbList = (JAXBSGEServerList)getStatus(
				"server", JAXBSGEServerList.class);

		if (jaxbList.getServer().size() == 0) {
			throw new SchedulerException("cannot get any server information");
		}

		// return first found server
		JAXBSGEServer server = jaxbList.getServer().get(0);
		return new SGEServerStatus(server);

	}

	public Collection<NodeStatus> getNodeStatus() throws SchedulerException {

		JAXBSGENodeList jaxbList = (JAXBSGENodeList)getStatus(
				"node", JAXBSGENodeList.class);

		Collection<NodeStatus> plusList = new LinkedList<NodeStatus>();
		for (JAXBSGENode node : jaxbList.getNode()) {
			String nodename = node.getEHName();
			if (nodename.equals("global") || nodename.equals("template")) {
				// ignore these node information
				continue;
			}
			plusList.add(new SGENodeStatus(node));
		}
		return plusList;

	}

	public Collection<QueueStatus> getQueueStatus() throws SchedulerException {

		JAXBSGEQueueList jaxbList = (JAXBSGEQueueList)getStatus(
				"queue", JAXBSGEQueueList.class);

		Collection<QueueStatus> plusList = new LinkedList<QueueStatus>();
		for (JAXBSGEQueue queue : jaxbList.getQueue()) {
			plusList.add(new SGEQueueStatus(queue));
		}
		return plusList;

	}

	public Collection<JobStatus> getJobStatus() throws SchedulerException {

		/*
		 * WE IGNORE 'queue' parameter.
		 */
		JAXBSGEJobList jaxbList = (JAXBSGEJobList)getStatus(
				"job", JAXBSGEJobList.class);

		Collection<JobStatus> plusList = new LinkedList<JobStatus>();
		for (JAXBSGEJob jaxbJob : jaxbList.getJob()) {
			SGEJobStatus sgeJob = new SGEJobStatus(jaxbJob);
			// add *tasks* to plusList, not job
			plusList.addAll(sgeJob.getAllTasks());
		}
		/*
		 * Return all *tasks* of all jobs in all queue always.
		 */
		return plusList;

	}

	public LinkedHashMap<String, SGEHostGroup> getHostGroups() throws SchedulerException {

		JAXBSGEHostgroupList jaxbList = (JAXBSGEHostgroupList)getStatus(
				"hgroup", JAXBSGEHostgroupList.class);

		LinkedHashMap<String, SGEHostGroup> plusMap =
			new LinkedHashMap<String, SGEHostGroup>();
		for (JAXBSGEHostgroup hgroup : jaxbList.getHgroup()) {
			plusMap.put(hgroup.getHGRPName(), new SGEHostGroup(hgroup));
		}
		return plusMap;

	}

	public LinkedHashMap<String, SGEUserset> getUsersets() throws SchedulerException {

		JAXBSGEUsersetList jaxbList = (JAXBSGEUsersetList)getStatus(
				"userset", JAXBSGEUsersetList.class);

		LinkedHashMap<String, SGEUserset> plusMap =
			new LinkedHashMap<String, SGEUserset>();
		for (USElement e : jaxbList.getUserset()) {
			plusMap.put(e.getUSName(), new SGEUserset(e));
		}
		return plusMap;

	}

	public LinkedHashMap<String, SGEProject> getProjects() throws SchedulerException {

		JAXBSGEProjectList jaxbList = (JAXBSGEProjectList)getStatus(
				"project", JAXBSGEProjectList.class);

		LinkedHashMap<String, SGEProject> plusMap =
			new LinkedHashMap<String, SGEProject>();
		for (JAXBSGEProject project : jaxbList.getProject()) {
			plusMap.put(project.getUPName(), new SGEProject(project));
		}
		return plusMap;

	}

	public LinkedHashMap<String, SGEParallelEnv> getParallelEnvs() throws SchedulerException {

		JAXBSGEPeList jaxbList = (JAXBSGEPeList)getStatus(
				"pe", JAXBSGEPeList.class);

		LinkedHashMap<String, SGEParallelEnv> plusMap =
			new LinkedHashMap<String, SGEParallelEnv>();
		for (JAXBSGEPe pe : jaxbList.getPe()) {
			plusMap.put(pe.getPEName(), new SGEParallelEnv(pe));
		}
		return plusMap;

	}

	private String checkReply() throws SchedulerException {

		/*
		 * reply maybe
		 * 	"OK"
		 * 	"OK some text (maybe new jobid etc)"
		 * 	"some error message"
		 */
		String reply = operatord.receiveReply();
		if (reply.startsWith("OK")) {
			if (reply.length() > 3) {
				return reply.substring(3);
			} else {
				return null;
			}
		} else {
			throw new SchedulerException(reply);
		}

	}

	private String sendOrder(String order) throws SchedulerException {

		operatord.sendRequest(order);
		return checkReply();

	}

	private void runJob(String jobId, String nodeNames) throws SchedulerException {

		sendOrder(String.format(Locale.US, "RunJob %s %s", jobId, nodeNames));

	}

	private SGEQInstance findQInsOnNode(NodeStatus node) throws SchedulerException {
		
		String nodeName = node.getName();
		Collection<QueueStatus> queues = getQueueStatus();
		for (QueueStatus q : queues) {
			SGEQueueStatus queue = (SGEQueueStatus)q;
			SGEQInstance qins = queue.getInstanceOnNode(nodeName);
			if (qins != null) {
				return qins;
			}
		}
		return null;
		
	}
	
	public void runJob(JobID jobID, NodeStatus node) throws SchedulerException {

		SGEQInstance qins = null;
		if (node instanceof SGEQInstance) {
			qins = (SGEQInstance)node;
		} else if (node instanceof SGENodeStatus) {
			qins = findQInsOnNode(node);
		}
		if (qins != null) {
			final int useSlot = 1;
			String arg = String.format("%s %d %d", 
					qins.getFullName(), qins.getVersion(), useSlot);
			runJob(jobID.toString(), arg);
		}
		
	}

	public void runJob(JobStatus job, Collection<NodeInfo> nodes)
			throws SchedulerException {

		assert(job instanceof SGETaskStatus);

		SGETaskStatus task = (SGETaskStatus)job;
		StringBuffer sb = new StringBuffer();
		for (NodeInfo node : nodes) {
			SGEQInstance ins = (SGEQInstance)node.getStatus();
			sb.append(String.format("%s %d %d ",
					node.getName(), ins.getVersion(),
					node.getLaunchProcessNum()));
		}

		String peName = task.getTargetParallelEnvName();
		if (TextUtil.isValid(peName)) {
			sb.append(peName);
		}

		runJob(task.getJobID().toString(), sb.toString());

		for (NodeInfo node : nodes) {
			SGEQInstance ins = (SGEQInstance)node.getStatus();
			ins.incrementUsedSlotNum(node.getLaunchProcessNum());
		}
	}

	private String operateJob(String order, JobID jobID) throws SchedulerException {

		return sendOrder(String.format(Locale.US, "%s %s",
				order, jobID.toString()));

	}

	public void deleteJob(JobID jobID) throws SchedulerException {

		// delete unfinished (waiting or running) job (shown by qstat)
		operateJob("DeleteJob", jobID);

	}

	public void expireJob(JobID jobID) throws SchedulerException {

		// delete finished job (cannot be seen by qstat)
		operateJob("ExpireJob", jobID);

	}

	public void suspendJob(JobID jobID) throws SchedulerException {

		operateJob("SuspendJob", jobID);

	}

	public void resumeJob(JobID jobID) throws SchedulerException {

		operateJob("ResumeJob", jobID);

	}

	public void holdJob(JobID jobID) throws SchedulerException {

		operateJob("HoldJob", jobID);

	}

	public void releaseJob(JobID jobID) throws SchedulerException {

		operateJob("ReleaseJob", jobID);

	}

	public void rerunJob(JobStatus job) throws SchedulerException {

		SGETaskStatus task = (SGETaskStatus)job;
		JobID jobID = task.getJobID();
		SGEJobID sgeID = (SGEJobID)jobID;
		String oldid = Integer.toString(sgeID.jobID());
		String newid = null;

		/*
		 * To Rerun job of 123.1, we must send
		 * 		"ResubmitJob 123 owner uid group gid"
		 * 		"DeleteJob 123"
		 * NOTE: SGE only allows resubmission by unit of a job, not a task.
		 * 'newid' doesn't contain task id. It maybe "124" etc.
		 */
		String owner = task.getOwner();
		String uid = Integer.toString(task.getUid());
		String group = task.getOwnersGroup();
		String gid = Integer.toString(task.getGid());
		newid = sendOrder("ResubmitJob " + oldid + " "
				+ owner + " " + uid + " " + group + " " + gid);

		try {
			sendOrder("DeleteJob " + oldid);
		} catch (SchedulerException e1) {
			try {
				sendOrder("DeleteJob " + newid);
			} catch (SchedulerException e2) {
				// recover failed, we cannot do anymore.
				// oldid & newid jobs remain
				throw new SchedulerException("failed to delete old job at RerunJob");
			}
		}

	}

	public void signalJob(JobID jobID, SignalType signal)
			throws SchedulerException {

		throw new SchedulerException("signalJob() for SGE is not supported!");

	}

}
